from langchain_chroma import Chroma
from langchain_community.document_loaders import TextLoader
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_text_splitters import CharacterTextSplitter

loader = TextLoader("/Users/wangliang/python/langchain/langchain-study/rag/元春省亲.txt")
documents = loader.load()
page_content = documents[0].page_content

text_splitter = CharacterTextSplitter.from_tiktoken_encoder(
    encoding_name="cl100k_base", chunk_size=1200, chunk_overlap=100
)
texts = text_splitter.split_text(page_content)

embeddings_model = OpenAIEmbeddings(base_url="https://api.openai-hk.com/v1",
                                    api_key="hk-0amgwp10000255022bdd816341db25b54dc2e46787aee69f",
                                    model="text-embedding-3-small")

vector = Chroma.from_texts(texts, embeddings_model)
retriever = vector.as_retriever()

docs = retriever.invoke("是谁让宝玉来园中玩耍的")
template = ChatPromptTemplate.from_messages([
    ('system', '根据我提供的下述信息，回答用户的问题。提供的信息：{context}'),
    ('user', '{input}')
])
prompt = template.invoke({"context": docs[0].page_content, "input": "是谁让宝玉来园中玩耍的"})
print(prompt)

model = ChatOpenAI(base_url="https://api.openai-hk.com/v1",
                   api_key="hk-0amgwp10000255022bdd816341db25b54dc2e46787aee69f",
                   model="gpt-4o-mini")
result = model.invoke(prompt)
print(result)
